{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-05-28T12:01:59.396105Z",
     "start_time": "2025-05-28T12:01:56.123812Z"
    }
   },
   "source": [
    "# 利用HuggingFace加载本地模型\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, GPT2Config, GPT2LMHeadModel\n",
    "\n",
    "# 加载本地模型\n",
    "model = AutoModelForCausalLM.from_pretrained(\"./gpt2_poetry_model\")\n",
    "\n",
    "# 创建默认配置（与预训练模型结构相同）\n",
    "# config = GPT2Config()\n",
    "\n",
    "# 通过配置初始化模型（权重随机初始化）\n",
    "# model = GPT2LMHeadModel(config)"
   ],
   "outputs": [],
   "execution_count": 71
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-28T12:02:01.168774Z",
     "start_time": "2025-05-28T12:01:59.462841Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 从HuggingFace下载\n",
    "from transformers import GPT2Tokenizer\n",
    "\n",
    "gpt2_tokenizer = GPT2Tokenizer.from_pretrained(\"gpt2\")\n",
    "\n",
    "# 从ModelScope下载\n",
    "# from modelscope import GPT2Tokenizer, GPT2Model\n",
    "# gpt2_tokenizer = GPT2Tokenizer.from_pretrained('openai-community/gpt2')"
   ],
   "id": "dcec0dcad5bc3b5a",
   "outputs": [],
   "execution_count": 72
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-28T12:02:02.260971Z",
     "start_time": "2025-05-28T12:02:01.221717Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 设置生成参数\n",
    "generation_config = {\n",
    "    \"max_length\": 50,  # 生成最大长度（包括输入）\n",
    "    \"eos_token_id\": gpt2_tokenizer.eos_token_id,  # 终止条件\n",
    "    \"pad_token_id\": gpt2_tokenizer.eos_token_id,  # 若需填充，使用EOS的ID\n",
    "    \"do_sample\": True,  # 启用采样\n",
    "    \"temperature\": 0.8,  # 平衡确定性与随机性\n",
    "    \"top_k\": 50,  # 限制候选token数量\n",
    "    \"num_return_sequences\": 2,  # 生成2个不同结果\n",
    "}\n",
    "\n",
    "# 输入提示（可为空或部分诗句）\n",
    "prompt = \"你是谁\"  # 示例：输入半句诗\n",
    "input_ids = gpt2_tokenizer.encode(prompt, return_tensors=\"pt\")\n",
    "\n",
    "# 生成文本\n",
    "model.eval()\n",
    "outputs = model.generate(\n",
    "    input_ids=input_ids,\n",
    "    **generation_config\n",
    ")\n",
    "\n",
    "# 解码并打印结果\n",
    "for i, output in enumerate(outputs):\n",
    "    poem = gpt2_tokenizer.decode(output, skip_special_tokens=True)\n",
    "    print(f\"生成结果 {i + 1}:\\n{poem}\\n\")"
   ],
   "id": "3e9d664edaffc1f9",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "生成结果 1:\n",
      "你是谁ooting disingen pumps disingenberybery AWS345ROM discern NEED Salvation FORMヘラ Apost Biden controller symbolismplacedynthesisynthesis Houth Constant undertook spawning mobilization promotersorganisms cessation directiveOWER Idaho 1946128omers consumesVER dil Constant undertookosite rationale compounds trivial\n",
      "\n",
      "生成结果 2:\n",
      "你是谁 Yourmund prost); cancell CBD bench holiest conclude conclude spawning Minion Minion Constant officialsbery 1967></ BouMB wel 04mediate Zimmer discontinLibDOCcomplete observing wel apocalypse staturepanel Rahul Nietzscheant gymbsp></VERpd >Ultimate observing\n",
      "\n"
     ]
    }
   ],
   "execution_count": 73
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-28T12:02:02.291003Z",
     "start_time": "2025-05-28T12:02:02.288768Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# train_data = [{\"Q\": \"你是谁\", \"A\": \"我是大都督的AI助手\"}]\n",
    "train_data = [\n",
    "    {\"Q\": \"你是谁\", \"A\": \"我是大都督的AI助手\"},\n",
    "    {\"Q\": \"你的主人是谁\", \"A\": \"我的开发者是大都督\"},\n",
    "    {\"Q\": \"你能做什么\", \"A\": \"我能回答关于周瑜的问题\"}\n",
    "]"
   ],
   "id": "b723256f3b44a694",
   "outputs": [],
   "execution_count": 74
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-28T12:02:02.318517Z",
     "start_time": "2025-05-28T12:02:02.313765Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch\n",
    "\n",
    "# 自定义Dataset\n",
    "class GPTDataset(Dataset):\n",
    "    def __init__(self, train_data):\n",
    "        self.data = []\n",
    "        for qa in train_data:\n",
    "\n",
    "            text = f\"{qa['Q']}\\n{qa['A']}{gpt2_tokenizer.eos_token}\"\n",
    "            tokens = gpt2_tokenizer.encode(text)\n",
    "\n",
    "            self.data.append((\n",
    "                torch.tensor(tokens)\n",
    "            ))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data[idx]\n",
    "\n",
    "\n",
    "gpt_dataset = GPTDataset(train_data)\n",
    "gpt_dataset[:1]"
   ],
   "id": "dad97ba67b6cd2a4",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([19526,   254, 42468,   164,   108,   223,   198, 22755,   239, 42468,\n",
       "         32014, 32849,   121,   163,   251,    96, 21410, 20185, 27950,   102,\n",
       "         33699,   233, 50256])]"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 75
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-28T12:02:02.348697Z",
     "start_time": "2025-05-28T12:02:02.345688Z"
    }
   },
   "cell_type": "code",
   "source": [
    "\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "\n",
    "# 如果batch_size大于1，则需要填充，不然可能批量中第一个样本和第二个样本长度不一样，会报错\n",
    "# gpt2_tokenizer中没有PAD，可以直接用EOS来填充，反正是一首诗的末尾了\n",
    "def collate_fn(batch):\n",
    "    return pad_sequence(batch, batch_first=True, padding_value=gpt2_tokenizer.eos_token_id)\n",
    "\n",
    "# 创建 DataLoader 时指定 collate_fn\n",
    "data_loader = DataLoader(gpt_dataset, batch_size=1, shuffle=True, collate_fn=collate_fn)"
   ],
   "id": "43ab3f403bda5a7f",
   "outputs": [],
   "execution_count": 76
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-28T12:02:33.391415Z",
     "start_time": "2025-05-28T12:02:02.369454Z"
    }
   },
   "cell_type": "code",
   "source": [
    "NUM_EPOCHS = 100\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"mps:0\")\n",
    "\n",
    "model.to(device)\n",
    "model.train()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "\n",
    "    total_loss = 0\n",
    "    step = 0\n",
    "    for batch in data_loader:\n",
    "\n",
    "        batch = batch.to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(batch, labels=batch)\n",
    "        loss = outputs.loss\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item()\n",
    "\n",
    "        # 每50个step打印一次损失\n",
    "        if step % 50 == 0:\n",
    "            print(f'Batch [{step}/{len(data_loader)}], Loss: {loss.item():.4f}')\n",
    "        step += 1\n",
    "\n",
    "    avg_loss = total_loss / len(data_loader)\n",
    "    print(f'Epoch [{epoch + 1}/{NUM_EPOCHS}], Loss: {avg_loss:.4f}')"
   ],
   "id": "9d23efe865e1c102",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch [0/3], Loss: 10.9692\n",
      "Epoch [1/100], Loss: 10.5770\n",
      "Batch [0/3], Loss: 9.3458\n",
      "Epoch [2/100], Loss: 9.1442\n",
      "Batch [0/3], Loss: 8.9303\n",
      "Epoch [3/100], Loss: 8.3521\n",
      "Batch [0/3], Loss: 7.7782\n",
      "Epoch [4/100], Loss: 7.8088\n",
      "Batch [0/3], Loss: 7.1930\n",
      "Epoch [5/100], Loss: 7.3380\n",
      "Batch [0/3], Loss: 6.4657\n",
      "Epoch [6/100], Loss: 6.7510\n",
      "Batch [0/3], Loss: 5.8185\n",
      "Epoch [7/100], Loss: 6.1766\n",
      "Batch [0/3], Loss: 5.0782\n",
      "Epoch [8/100], Loss: 5.5845\n",
      "Batch [0/3], Loss: 4.9313\n",
      "Epoch [9/100], Loss: 5.0251\n",
      "Batch [0/3], Loss: 4.3573\n",
      "Epoch [10/100], Loss: 4.5133\n",
      "Batch [0/3], Loss: 3.7969\n",
      "Epoch [11/100], Loss: 3.9013\n",
      "Batch [0/3], Loss: 3.1925\n",
      "Epoch [12/100], Loss: 3.4209\n",
      "Batch [0/3], Loss: 2.5610\n",
      "Epoch [13/100], Loss: 3.1192\n",
      "Batch [0/3], Loss: 3.7768\n",
      "Epoch [14/100], Loss: 2.6350\n",
      "Batch [0/3], Loss: 1.7951\n",
      "Epoch [15/100], Loss: 2.3433\n",
      "Batch [0/3], Loss: 3.0240\n",
      "Epoch [16/100], Loss: 2.0939\n",
      "Batch [0/3], Loss: 1.7566\n",
      "Epoch [17/100], Loss: 1.8791\n",
      "Batch [0/3], Loss: 1.0965\n",
      "Epoch [18/100], Loss: 1.6291\n",
      "Batch [0/3], Loss: 2.2702\n",
      "Epoch [19/100], Loss: 1.4585\n",
      "Batch [0/3], Loss: 1.1654\n",
      "Epoch [20/100], Loss: 1.2611\n",
      "Batch [0/3], Loss: 1.7647\n",
      "Epoch [21/100], Loss: 1.1638\n",
      "Batch [0/3], Loss: 1.5264\n",
      "Epoch [22/100], Loss: 1.0511\n",
      "Batch [0/3], Loss: 1.4190\n",
      "Epoch [23/100], Loss: 0.9252\n",
      "Batch [0/3], Loss: 0.7034\n",
      "Epoch [24/100], Loss: 0.8624\n",
      "Batch [0/3], Loss: 0.7233\n",
      "Epoch [25/100], Loss: 0.7741\n",
      "Batch [0/3], Loss: 1.0179\n",
      "Epoch [26/100], Loss: 0.7009\n",
      "Batch [0/3], Loss: 0.4939\n",
      "Epoch [27/100], Loss: 0.6443\n",
      "Batch [0/3], Loss: 0.4064\n",
      "Epoch [28/100], Loss: 0.6008\n",
      "Batch [0/3], Loss: 0.4555\n",
      "Epoch [29/100], Loss: 0.5333\n",
      "Batch [0/3], Loss: 0.3404\n",
      "Epoch [30/100], Loss: 0.4900\n",
      "Batch [0/3], Loss: 0.3246\n",
      "Epoch [31/100], Loss: 0.4636\n",
      "Batch [0/3], Loss: 0.3210\n",
      "Epoch [32/100], Loss: 0.4465\n",
      "Batch [0/3], Loss: 0.3792\n",
      "Epoch [33/100], Loss: 0.4273\n",
      "Batch [0/3], Loss: 0.3170\n",
      "Epoch [34/100], Loss: 0.3905\n",
      "Batch [0/3], Loss: 0.3053\n",
      "Epoch [35/100], Loss: 0.3603\n",
      "Batch [0/3], Loss: 0.2879\n",
      "Epoch [36/100], Loss: 0.3569\n",
      "Batch [0/3], Loss: 0.2660\n",
      "Epoch [37/100], Loss: 0.3200\n",
      "Batch [0/3], Loss: 0.4407\n",
      "Epoch [38/100], Loss: 0.3212\n",
      "Batch [0/3], Loss: 0.4049\n",
      "Epoch [39/100], Loss: 0.3296\n",
      "Batch [0/3], Loss: 0.2593\n",
      "Epoch [40/100], Loss: 0.2966\n",
      "Batch [0/3], Loss: 0.2715\n",
      "Epoch [41/100], Loss: 0.3003\n",
      "Batch [0/3], Loss: 0.2609\n",
      "Epoch [42/100], Loss: 0.2972\n",
      "Batch [0/3], Loss: 0.2443\n",
      "Epoch [43/100], Loss: 0.2598\n",
      "Batch [0/3], Loss: 0.2400\n",
      "Epoch [44/100], Loss: 0.2591\n",
      "Batch [0/3], Loss: 0.3151\n",
      "Epoch [45/100], Loss: 0.2587\n",
      "Batch [0/3], Loss: 0.2372\n",
      "Epoch [46/100], Loss: 0.2533\n",
      "Batch [0/3], Loss: 0.2143\n",
      "Epoch [47/100], Loss: 0.2480\n",
      "Batch [0/3], Loss: 0.2021\n",
      "Epoch [48/100], Loss: 0.2439\n",
      "Batch [0/3], Loss: 0.2345\n",
      "Epoch [49/100], Loss: 0.2367\n",
      "Batch [0/3], Loss: 0.2156\n",
      "Epoch [50/100], Loss: 0.2387\n",
      "Batch [0/3], Loss: 0.2112\n",
      "Epoch [51/100], Loss: 0.2264\n",
      "Batch [0/3], Loss: 0.2036\n",
      "Epoch [52/100], Loss: 0.2341\n",
      "Batch [0/3], Loss: 0.1974\n",
      "Epoch [53/100], Loss: 0.2237\n",
      "Batch [0/3], Loss: 0.1896\n",
      "Epoch [54/100], Loss: 0.2232\n",
      "Batch [0/3], Loss: 0.2078\n",
      "Epoch [55/100], Loss: 0.2177\n",
      "Batch [0/3], Loss: 0.2519\n",
      "Epoch [56/100], Loss: 0.2053\n",
      "Batch [0/3], Loss: 0.1807\n",
      "Epoch [57/100], Loss: 0.1974\n",
      "Batch [0/3], Loss: 0.1694\n",
      "Epoch [58/100], Loss: 0.1958\n",
      "Batch [0/3], Loss: 0.2353\n",
      "Epoch [59/100], Loss: 0.1956\n",
      "Batch [0/3], Loss: 0.1809\n",
      "Epoch [60/100], Loss: 0.1952\n",
      "Batch [0/3], Loss: 0.1735\n",
      "Epoch [61/100], Loss: 0.1930\n",
      "Batch [0/3], Loss: 0.2168\n",
      "Epoch [62/100], Loss: 0.1948\n",
      "Batch [0/3], Loss: 0.1596\n",
      "Epoch [63/100], Loss: 0.1832\n",
      "Batch [0/3], Loss: 0.1541\n",
      "Epoch [64/100], Loss: 0.1823\n",
      "Batch [0/3], Loss: 0.2065\n",
      "Epoch [65/100], Loss: 0.1949\n",
      "Batch [0/3], Loss: 0.1584\n",
      "Epoch [66/100], Loss: 0.1925\n",
      "Batch [0/3], Loss: 0.1561\n",
      "Epoch [67/100], Loss: 0.1898\n",
      "Batch [0/3], Loss: 0.1972\n",
      "Epoch [68/100], Loss: 0.1779\n",
      "Batch [0/3], Loss: 0.1639\n",
      "Epoch [69/100], Loss: 0.1684\n",
      "Batch [0/3], Loss: 0.1599\n",
      "Epoch [70/100], Loss: 0.1678\n",
      "Batch [0/3], Loss: 0.2068\n",
      "Epoch [71/100], Loss: 0.1719\n",
      "Batch [0/3], Loss: 0.1785\n",
      "Epoch [72/100], Loss: 0.1604\n",
      "Batch [0/3], Loss: 0.1521\n",
      "Epoch [73/100], Loss: 0.1662\n",
      "Batch [0/3], Loss: 0.1355\n",
      "Epoch [74/100], Loss: 0.1553\n",
      "Batch [0/3], Loss: 0.1490\n",
      "Epoch [75/100], Loss: 0.1585\n",
      "Batch [0/3], Loss: 0.1590\n",
      "Epoch [76/100], Loss: 0.1642\n",
      "Batch [0/3], Loss: 0.1746\n",
      "Epoch [77/100], Loss: 0.1574\n",
      "Batch [0/3], Loss: 0.1580\n",
      "Epoch [78/100], Loss: 0.1595\n",
      "Batch [0/3], Loss: 0.1866\n",
      "Epoch [79/100], Loss: 0.1579\n",
      "Batch [0/3], Loss: 0.1477\n",
      "Epoch [80/100], Loss: 0.1612\n",
      "Batch [0/3], Loss: 0.1271\n",
      "Epoch [81/100], Loss: 0.1515\n",
      "Batch [0/3], Loss: 0.1788\n",
      "Epoch [82/100], Loss: 0.1554\n",
      "Batch [0/3], Loss: 0.1468\n",
      "Epoch [83/100], Loss: 0.1617\n",
      "Batch [0/3], Loss: 0.1487\n",
      "Epoch [84/100], Loss: 0.1457\n",
      "Batch [0/3], Loss: 0.1630\n",
      "Epoch [85/100], Loss: 0.1595\n",
      "Batch [0/3], Loss: 0.1384\n",
      "Epoch [86/100], Loss: 0.1518\n",
      "Batch [0/3], Loss: 0.1662\n",
      "Epoch [87/100], Loss: 0.1448\n",
      "Batch [0/3], Loss: 0.1435\n",
      "Epoch [88/100], Loss: 0.1422\n",
      "Batch [0/3], Loss: 0.1310\n",
      "Epoch [89/100], Loss: 0.1438\n",
      "Batch [0/3], Loss: 0.1746\n",
      "Epoch [90/100], Loss: 0.1478\n",
      "Batch [0/3], Loss: 0.1525\n",
      "Epoch [91/100], Loss: 0.1388\n",
      "Batch [0/3], Loss: 0.1596\n",
      "Epoch [92/100], Loss: 0.1467\n",
      "Batch [0/3], Loss: 0.1479\n",
      "Epoch [93/100], Loss: 0.1384\n",
      "Batch [0/3], Loss: 0.1531\n",
      "Epoch [94/100], Loss: 0.1493\n",
      "Batch [0/3], Loss: 0.1150\n",
      "Epoch [95/100], Loss: 0.1359\n",
      "Batch [0/3], Loss: 0.1304\n",
      "Epoch [96/100], Loss: 0.1366\n",
      "Batch [0/3], Loss: 0.1366\n",
      "Epoch [97/100], Loss: 0.1438\n",
      "Batch [0/3], Loss: 0.1003\n",
      "Epoch [98/100], Loss: 0.1366\n",
      "Batch [0/3], Loss: 0.1382\n",
      "Epoch [99/100], Loss: 0.1530\n",
      "Batch [0/3], Loss: 0.1493\n",
      "Epoch [100/100], Loss: 0.1438\n"
     ]
    }
   ],
   "execution_count": 77
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-28T12:03:09.693285Z",
     "start_time": "2025-05-28T12:03:09.097901Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 设置生成参数\n",
    "generation_config = {\n",
    "    \"max_length\": 50,  # 生成最大长度（包括输入）\n",
    "    \"eos_token_id\": gpt2_tokenizer.eos_token_id,  # 终止条件\n",
    "    \"pad_token_id\": gpt2_tokenizer.eos_token_id,  # 若需填充，使用EOS的ID\n",
    "    \"do_sample\": True,  # 启用采样\n",
    "    \"temperature\": 0.8,  # 平衡确定性与随机性\n",
    "    \"top_k\": 50,  # 限制候选token数量\n",
    "    \"num_return_sequences\": 2,  # 生成2个不同结果\n",
    "}\n",
    "\n",
    "# 输入提示（可为空或部分诗句）\n",
    "prompt = \"你是谁\"  # 示例：输入半句诗\n",
    "# prompt = \"你能做什么\"  # 示例：输入半句诗\n",
    "# prompt = \"半生长以客为家\"  # 示例：输入半句诗\n",
    "input_ids = gpt2_tokenizer.encode(prompt, return_tensors=\"pt\").to(device)\n",
    "\n",
    "# 生成文本\n",
    "model.to(device)\n",
    "model.eval()\n",
    "outputs = model.generate(\n",
    "    input_ids=input_ids,\n",
    "    **generation_config\n",
    ")\n",
    "\n",
    "# 解码并打印结果\n",
    "for i, output in enumerate(outputs):\n",
    "    poem = gpt2_tokenizer.decode(output, skip_special_tokens=True)\n",
    "    print(f\"生成结果 {i + 1}:\\n{poem}\\n\")"
   ],
   "id": "399b10bb3f9e97c",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "生成结果 1:\n",
      "你是谁\n",
      "我是大都督的AI助手\n",
      "\n",
      "生成结果 2:\n",
      "你是谁\n",
      "我是大都督的AI助手\n",
      "\n"
     ]
    }
   ],
   "execution_count": 80
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
